#
# BSD 3-Clause License
#
# This file is part of the Basalt project.
# https://gitlab.com/VladyslavUsenko/basalt.git
#
# Copyright (c) 2022, Collabora Ltd.
# All rights reserved.
#
# Author: Mateo de Mayo <mateo.demayo@collabora.com>
#
# Modified for fisheye624 camera model by:
# Copyright (c) 2025, yanghe <yanghe215@gmail.com>
# 

"""
Computes jacobian expressions for the fisheye624 camera model.
"""
from itertools import count
from sympy import symbols, diff, simplify, cse, sqrt, atan
from sympy.printing import cxxcode
from sympy.codegen.rewriting import create_expand_pow_optimization

fx, fy, cx, cy, k1, k2, k3, k4, k5, k6, p1, p2, s1, s2, s3, s4 = symbols(
    "fx fy cx cy k1 k2 k3 k4 k5 k6 p1 p2 s1 s2 s3 s4"
)


def distort_symbolic(xp, yp):
    # compute r^2 and r
    rp2 = xp * xp + yp * yp
    rp = sqrt(rp2)
    
    # compute theta angle
    th = atan(rp)
    th2 = th * th
    
    # compute cos_phi and sin_phi
    cos_phi = xp / rp
    sin_phi = yp / rp
    
    # compute radial distortion polynomial
    theta_dist = th * (1 + th2 * (k1 + th2 * (k2 + th2 * (k3 + th2 * (k4 + th2 * (k5 + th2 * k6))))))
    
    # compute distorted coordinates
    xr = theta_dist * cos_phi
    yr = theta_dist * sin_phi
    rd2 = xr * xr + yr * yr
    
    # compute tangential distortion
    dx_tan = (2 * xr * xr + rd2) * p1 + 2 * xr * yr * p2
    dy_tan = (2 * yr * yr + rd2) * p2 + 2 * xr * yr * p1
    
    # compute prism distortion
    dx_prism = s1 * rd2 + s2 * rd2 * rd2
    dy_prism = s3 * rd2 + s4 * rd2 * rd2
    
    # compute final distorted coordinates
    xpp = xr + dx_tan + dx_prism
    ypp = yr + dy_tan + dy_prism
    
    return xpp, ypp


def project_symbolic(x, y, z):
    # compute normalized coordinates
    xp = x / z
    yp = y / z
    
    # apply distortion
    xpp, ypp = distort_symbolic(xp, yp)
    
    # apply intrinsics
    u = fx * xpp + cx
    v = fy * ypp + cy
    
    return u, v


def cse_as_cxx(cse_tuple, expr_names):
    expand_pow = create_expand_pow_optimization(20) # create_expand_pow_optimization() create a optimization function for expanding power operations
    betterformat = lambda e: cxxcode(expand_pow(e)) # convert expression to C++ code

    subexprs, rhss = cse_tuple # convert expression to C++ code

    res = ""
    for var, expr in subexprs:
        res += f"const Scalar {var} = {betterformat(expr)};\n"
    res += "\n"

    for lhs, rhs in zip(expr_names, rhss):
        res += f"const Scalar {lhs} = {betterformat(rhs)};\n"

    return res


def distort_jacobian_to_xpyp():
    # Requires uncommenting line above for xp, yp
    var_names = (symbols(f"v{i}") for i in count())
    xp, yp = symbols("xp yp")
    xpp, ypp = distort_symbolic(xp, yp)

    dxpp_dxp = simplify(diff(xpp, xp))
    dxpp_dyp = simplify(diff(xpp, yp))
    dypp_dxp = simplify(diff(ypp, xp))
    dypp_dyp = simplify(diff(ypp, yp))
    cse_tuple = cse([dxpp_dxp, dxpp_dyp, dypp_dxp, dypp_dyp], var_names)
    names = ["dxpp_dxp", "dxpp_dyp", "dypp_dxp", "dypp_dyp"]
    return cse_tuple, names


def project_jacobian_to_3dpoint():
    var_names = (symbols(f"v{i}") for i in count())
    x, y, z = symbols("x y z")
    u, v = project_symbolic(x, y, z)

    du_dx = simplify(diff(u, x))
    du_dy = simplify(diff(u, y))
    du_dz = simplify(diff(u, z))
    dv_dx = simplify(diff(v, x))
    dv_dy = simplify(diff(v, y))
    dv_dz = simplify(diff(v, z))
    cse_tuple = cse([du_dx, du_dy, du_dz, dv_dx, dv_dy, dv_dz], var_names)
    names = ["du_dx", "du_dy", "du_dz", "dv_dx", "dv_dy", "dv_dz"]
    return cse_tuple, names


def project_jacobian_to_intrinsics():
    var_names = (symbols(f"w{i}") for i in count())
    x, y, z = symbols("x y z")
    u, v = project_symbolic(x, y, z)

    du_fx = simplify(diff(u, fx))
    du_fy = simplify(diff(u, fy))
    du_cx = simplify(diff(u, cx))
    du_cy = simplify(diff(u, cy))
    du_k1 = simplify(diff(u, k1))
    du_k2 = simplify(diff(u, k2))
    du_k3 = simplify(diff(u, k3))
    du_k4 = simplify(diff(u, k4))
    du_k5 = simplify(diff(u, k5))
    du_k6 = simplify(diff(u, k6))
    du_p1 = simplify(diff(u, p1))
    du_p2 = simplify(diff(u, p2))
    du_s1 = simplify(diff(u, s1))
    du_s2 = simplify(diff(u, s2))
    du_s3 = simplify(diff(u, s3))
    du_s4 = simplify(diff(u, s4))

    dv_fx = simplify(diff(v, fx))
    dv_fy = simplify(diff(v, fy))
    dv_cx = simplify(diff(v, cx))
    dv_cy = simplify(diff(v, cy))
    dv_k1 = simplify(diff(v, k1))
    dv_k2 = simplify(diff(v, k2))
    dv_k3 = simplify(diff(v, k3))
    dv_k4 = simplify(diff(v, k4))
    dv_k5 = simplify(diff(v, k5))
    dv_k6 = simplify(diff(v, k6))
    dv_p1 = simplify(diff(v, p1))
    dv_p2 = simplify(diff(v, p2))
    dv_s1 = simplify(diff(v, s1))
    dv_s2 = simplify(diff(v, s2))
    dv_s3 = simplify(diff(v, s3))
    dv_s4 = simplify(diff(v, s4))

    # fmt:off
    cse_tuple = cse(
        [
            du_fx, du_fy, du_cx, du_cy, du_k1, du_k2, du_k3, du_k4, du_k5, du_k6,
            du_p1, du_p2, du_s1, du_s2, du_s3, du_s4,
            dv_fx, dv_fy, dv_cx, dv_cy, dv_k1, dv_k2, dv_k3, dv_k4, dv_k5, dv_k6,
            dv_p1, dv_p2, dv_s1, dv_s2, dv_s3, dv_s4,
        ],
        var_names
    )
    names = [
        "du_fx", "du_fy", "du_cx", "du_cy", "du_k1", "du_k2", "du_k3", "du_k4", "du_k5", "du_k6",
        "du_p1", "du_p2", "du_s1", "du_s2", "du_s3", "du_s4",
        "dv_fx", "dv_fy", "dv_cx", "dv_cy", "dv_k1", "dv_k2", "dv_k3", "dv_k4", "dv_k5", "dv_k6",
        "dv_p1", "dv_p2", "dv_s1", "dv_s2", "dv_s3", "dv_s4",
    ]
    # fmt:on

    return cse_tuple, names


def main():
    print("[Project jacobians of 2D point w.r.t 3D point]")
    res = project_jacobian_to_3dpoint()
    print(cse_as_cxx(*res))

    print("[Project jacobians of 2D point w.r.t intrinsics]")
    res = project_jacobian_to_intrinsics()
    print(cse_as_cxx(*res))

    print("[Distort jacobians of distorted 2D point from undistorted 2D point]")
    res = distort_jacobian_to_xpyp()
    print(cse_as_cxx(*res))

    print("[Unproject jacobians of 3D point w.r.t 2D point]")
    print("Not implemented yet")

    print("[Unproject jacobians of 3D point w.r.t intrinsics]")
    print("Not implemented yet")


if __name__ == "__main__":
    main()


'''
[Project jacobians of 2D point w.r.t 3D point]
const Scalar v0 = z*z;
const Scalar v1 = x*x;
const Scalar v2 = y*y;
const Scalar v3 = v1 + v2;
const Scalar v4 = v3*1.0/v0;
const Scalar v5 = std::pow(v4, 3.0/2.0);
const Scalar v6 = z*z*z;
const Scalar v7 = v5*v6;
const Scalar v8 = v0 + v3;
const Scalar v9 = v1*v8;
const Scalar v10 = 4*v9;
const Scalar v11 = std::sqrt(v4);
const Scalar v12 = std::atan(v11);
const Scalar v13 = v12*v12;
const Scalar v14 = k6*v13;
const Scalar v15 = k4 + v13*(k5 + v14);
const Scalar v16 = k3 + v13*v15;
const Scalar v17 = k2 + v13*v16;
const Scalar v18 = k1 + v13*v17;
const Scalar v19 = v13*v18 + 1;
const Scalar v20 = v13*(v19*v19);
const Scalar v21 = p2*v20;
const Scalar v22 = v21*y;
const Scalar v23 = 1.0/v11;
const Scalar v24 = v3*v3*v3;
const Scalar v25 = 2*v23*v24*v8*1.0/z;
const Scalar v26 = v13*(v13*(v13*(v13*(k5 + 2*v14) + v15) + v16) + v17) + v18;
const Scalar v27 = v13*v26;
const Scalar v28 = 2*v27;
const Scalar v29 = v19 + v28;
const Scalar v30 = v24*v29;
const Scalar v31 = v23*v30;
const Scalar v32 = 4*z;
const Scalar v33 = v19*v19*v19;
const Scalar v34 = v12*v12*v12;
const Scalar v35 = s2*v33*v34;
const Scalar v36 = v3*v8;
const Scalar v37 = v19*y;
const Scalar v38 = y*z;
const Scalar v39 = 8*v27;
const Scalar v40 = v3*v3;
const Scalar v41 = v12*v19;
const Scalar v42 = v40*v41;
const Scalar v43 = v29*v40;
const Scalar v44 = v23*v43;
const Scalar v45 = s1*v44;
const Scalar v46 = v36*v41;
const Scalar v47 = 3*v46;
const Scalar v48 = 3*v1 + v2;
const Scalar v49 = v0*v11;
const Scalar v50 = v29*v49;
const Scalar v51 = v41*v8;
const Scalar v52 = v48*v50 - v48*v51;
const Scalar v53 = 2*v41;
const Scalar v54 = v53*x;
const Scalar v55 = 1.0/v8;
const Scalar v56 = v55*1.0/v40;
const Scalar v57 = fx*v56;
const Scalar v58 = 1.0/v5*1.0/v6;
const Scalar v59 = 4*x;
const Scalar v60 = p2*v59;
const Scalar v61 = v2*v8;
const Scalar v62 = v20*v61;
const Scalar v63 = 2*v36;
const Scalar v64 = v38*x;
const Scalar v65 = v11*v41*v64;
const Scalar v66 = v29*v3*v64;
const Scalar v67 = v19*x;
const Scalar v68 = 8*v26*v34*v49;
const Scalar v69 = 4*v35;
const Scalar v70 = v53*y;
const Scalar v71 = z*z*z*z*z*z*z*z*z*z*z*z*z*z*z*z*z;
const Scalar v72 = std::pow(v4, 17.0/2.0);
const Scalar v73 = v71*v72;
const Scalar v74 = v28*v73;
const Scalar v75 = v29*v41*(v3*v3*v3*v3*v3*v3*v3*v3);
const Scalar v76 = -v0*v5*v51 + v11*v46 + v43;
const Scalar v77 = v76*(v3*v3*v3*v3*v3*v3*v3);
const Scalar v78 = v53*v76*(v3*v3*v3*v3*v3*v3);
const Scalar v79 = v55*1.0/v71*1.0/v72;
const Scalar v80 = p1*y;
const Scalar v81 = v20*v80;
const Scalar v82 = 4*s4*v33*v34;
const Scalar v83 = s3*v44;
const Scalar v84 = v1 + 3*v2;
const Scalar v85 = v50*v84 - v51*v84;
const Scalar v86 = fy*v56;
const Scalar v87 = p1*x;

const Scalar du_dx = v57*v58*(v1*v31 - v10*v22*v7 + v22*v25 + v30*v32*v35*x + v42*(p2*v1*v32*v37 + p2*v1*v38*v39 + v36 - v9) + v54*v7*(p1*(v47 + v52) + v45));
const Scalar du_dy = v57*(p2*v2*v67*v68 + v21*v63*x + v44*v69*y - v60*v62 + v65*(4*p2*v19*y*z - v8) + v66 + v70*(p1*(v46 + v52) + v45));
const Scalar du_dz = -fx*v79*(v60*v75*y + v67*v73 + v69*v77 + v74*x + v78*(p1*v48 + s1*v3));
const Scalar dv_dx = v86*(p1*v1*v37*v68 - v10*v81 + v44*v82*x + v54*(p2*(v46 + v85) + v83) + v63*v81 + v65*(4*p1*v19*x*z - v8) + v66);
const Scalar dv_dy = v58*v86*(-p1*v59*v62*v7 + v2*v31 + v20*v25*v87 + v30*v38*v82 + v42*(p1*v2*v32*v67 + v2*v39*v87*z + v36 - v61) + v7*v70*(p2*(v47 + v85) + v83));
const Scalar dv_dz = -fy*v79*(v37*v73 + v59*v75*v80 + v74*y + v77*v82 + v78*(p2*v84 + s3*v3));

[Project jacobians of 2D point w.r.t intrinsics]
const Scalar w0 = x*x;
const Scalar w1 = y*y;
const Scalar w2 = w0 + w1;
const Scalar w3 = std::sqrt(w2*1.0/(z*z));
const Scalar w4 = w3*z;
const Scalar w5 = w4*x;
const Scalar w6 = std::atan(w3);
const Scalar w7 = w6*w6;
const Scalar w8 = w7*(k1 + w7*(k2 + w7*(k3 + w7*(k4 + w7*(k5 + k6*w7))))) + 1;
const Scalar w9 = w6*w6*w6;
const Scalar w10 = w2*w9*(w8*w8*w8);
const Scalar w11 = s2*w10;
const Scalar w12 = 2*x*y;
const Scalar w13 = 3*w0 + w1;
const Scalar w14 = w6*w8;
const Scalar w15 = w14*(p1*w13 + p2*w12 + s1*w2);
const Scalar w16 = 1.0/w2;
const Scalar w17 = w14*w16;
const Scalar w18 = w16*w9;
const Scalar w19 = fx*(4*w11 + 2*w15 + w5);
const Scalar w20 = w6*w6*w6*w6*w6;
const Scalar w21 = w16*w19;
const Scalar w22 = w6*w6*w6*w6*w6*w6*w6;
const Scalar w23 = w6*w6*w6*w6*w6*w6*w6*w6*w6;
const Scalar w24 = w6*w6*w6*w6*w6*w6*w6*w6*w6*w6*w6;
const Scalar w25 = w6*w6*w6*w6*w6*w6*w6*w6*w6*w6*w6*w6*w6;
const Scalar w26 = w7*(w8*w8);
const Scalar w27 = fx*w26;
const Scalar w28 = w16*w27;
const Scalar w29 = (w6*w6*w6*w6)*(w8*w8*w8*w8);
const Scalar w30 = w4*y;
const Scalar w31 = s4*w10;
const Scalar w32 = w0 + 3*w1;
const Scalar w33 = w14*(p1*w12 + p2*w32 + s3*w2);
const Scalar w34 = fy*(w30 + 4*w31 + 2*w33);
const Scalar w35 = w16*w34;
const Scalar w36 = fy*w26;
const Scalar w37 = w16*w36;

const Scalar du_fx = w17*(w11 + w15 + w5);
const Scalar du_fy = 0;
const Scalar du_cx = 1;
const Scalar du_cy = 0;
const Scalar du_k1 = w18*w19;
const Scalar du_k2 = w20*w21;
const Scalar du_k3 = w21*w22;
const Scalar du_k4 = w21*w23;
const Scalar du_k5 = w21*w24;
const Scalar du_k6 = w21*w25;
const Scalar du_p1 = w13*w28;
const Scalar du_p2 = w12*w28;
const Scalar du_s1 = w27;
const Scalar du_s2 = fx*w29;
const Scalar du_s3 = 0;
const Scalar du_s4 = 0;
const Scalar dv_fx = 0;
const Scalar dv_fy = w17*(w30 + w31 + w33);
const Scalar dv_cx = 0;
const Scalar dv_cy = 1;
const Scalar dv_k1 = w18*w34;
const Scalar dv_k2 = w20*w35;
const Scalar dv_k3 = w22*w35;
const Scalar dv_k4 = w23*w35;
const Scalar dv_k5 = w24*w35;
const Scalar dv_k6 = w25*w35;
const Scalar dv_p1 = w12*w37;
const Scalar dv_p2 = w32*w37;
const Scalar dv_s1 = 0;
const Scalar dv_s2 = 0;
const Scalar dv_s3 = w36;
const Scalar dv_s4 = fy*w29;

[Distort jacobians of distorted 2D point from undistorted 2D point]
const Scalar v0 = xp*xp;
const Scalar v1 = yp*yp;
const Scalar v2 = v0 + v1;
const Scalar v3 = v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2;
const Scalar v4 = std::atan(std::sqrt(v2));
const Scalar v5 = v4*v4;
const Scalar v6 = k6*v5;
const Scalar v7 = k4 + v5*(k5 + v6);
const Scalar v8 = k3 + v5*v7;
const Scalar v9 = k2 + v5*v8;
const Scalar v10 = k1 + v5*v9;
const Scalar v11 = v10*v5 + 1;
const Scalar v12 = v2 + 1;
const Scalar v13 = v12*v4;
const Scalar v14 = v11*v13;
const Scalar v15 = v14*v3;
const Scalar v16 = v0*v11;
const Scalar v17 = v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2;
const Scalar v18 = v13*v17;
const Scalar v19 = std::pow(v2, 33.0/2.0);
const Scalar v20 = 2*yp;
const Scalar v21 = v12*v5*(v11*v11);
const Scalar v22 = p2*v21;
const Scalar v23 = std::pow(v2, 31.0/2.0);
const Scalar v24 = 4*yp;
const Scalar v25 = v23*v24;
const Scalar v26 = v5*(v10 + v5*(v5*(v5*(v5*(k5 + 2*v6) + v7) + v8) + v9));
const Scalar v27 = 2*v26;
const Scalar v28 = v11 + v27;
const Scalar v29 = v19*v28;
const Scalar v30 = v17*v28;
const Scalar v31 = p2*v4;
const Scalar v32 = 4*xp;
const Scalar v33 = v1*v11;
const Scalar v34 = v1*v27 + v33;
const Scalar v35 = v0*v27 + v16;
const Scalar v36 = v34 + v35;
const Scalar v37 = v11*v11*v11;
const Scalar v38 = v4*v4*v4;
const Scalar v39 = s2*v36*v37*v38;
const Scalar v40 = v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2;
const Scalar v41 = v2*v2*v2;
const Scalar v42 = v36*v41;
const Scalar v43 = s1*v42;
const Scalar v44 = v14*std::pow(v2, 7.0/2.0);
const Scalar v45 = 3*v44;
const Scalar v46 = 3*v0;
const Scalar v47 = 6*v26;
const Scalar v48 = v14*std::pow(v2, 5.0/2.0);
const Scalar v49 = v41*(v0*v47 + v11*v46 + v34) - v48*(v1 + v46);
const Scalar v50 = 2*xp;
const Scalar v51 = v11*v4;
const Scalar v52 = v50*v51;
const Scalar v53 = 1.0/v12;
const Scalar v54 = v53/std::pow(v2, 35.0/2.0);
const Scalar v55 = std::pow(v2, 25.0/2.0);
const Scalar v56 = v20*v51;
const Scalar v57 = v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2*v2;
const Scalar v58 = v23*v32;
const Scalar v59 = xp*yp;
const Scalar v60 = -v14*v23*v59 + v30*v59;
const Scalar v61 = v53*1.0/v3;
const Scalar v62 = s3*v42;
const Scalar v63 = v41*(v1*v47 + 3*v33 + v35) - v48*(v0 + 3*v1);
const Scalar v64 = p1*v21;
const Scalar v65 = s4*v36*v37*v38;
const Scalar v66 = p1*v4;

const Scalar dxpp_dxp = v54*(-v0*v22*v25 + v0*v29 + v15 - v16*v18 + v16*v24*v30*v31 + v17*v32*v39 + v19*v20*v22 + v40*v52*(p1*(v45 + v49) + v43));
const Scalar dxpp_dyp = v61*(-v1*v22*v32*v57 + v17*v22*v50 + v25*v39 + v28*v31*v33*v58 + v55*v56*(p1*(v44 + v49) + v43) + v60);
const Scalar dypp_dxp = v61*(-v0*v24*v57*v64 + v16*v25*v28*v66 + v17*v20*v64 + v52*v55*(p2*(v44 + v63) + v62) + v58*v65 + v60);
const Scalar dypp_dyp = v54*(v1*v29 - v1*v58*v64 + v15 + v17*v24*v65 - v18*v33 + v19*v50*v64 + v30*v32*v33*v66 + v40*v56*(p2*(v45 + v63) + v62));

[Unproject jacobians of 3D point w.r.t 2D point]
Not implemented yet
[Unproject jacobians of 3D point w.r.t intrinsics]
Not implemented yet
'''